#ifndef FPAL_HPP
#define FPAL_HPP

#include "ap_fixed.h" 
#include "hls_math.h" 
#include <iostream>   
#include <limits>     
#include <type_traits> 
#include <algorithm>  
#include <cmath>

// -----------------------------------------------------------
// 封装在命名空间 fpal 中
// -----------------------------------------------------------
namespace fpal {

// --- 1. 定点数类型定义 ---
typedef ap_fixed<24, 14, AP_RND_CONV, AP_SAT> fp_coord_t;
typedef ap_fixed<16, 5, AP_RND_CONV, AP_SAT> fp_vec_unit_t;
typedef ap_uint<16> fp_z_buffer_t;
typedef ap_uint<8> fp_color_channel_t;
typedef ap_fixed<17, 9, AP_RND_CONV, AP_SAT> fp_color_scale_t;


// --- 辅助结构体：用于将定点数类型打包成向量或颜色 ---
struct vec3_fp_coord {
    fp_coord_t x;
    fp_coord_t y;
    fp_coord_t z;
};
struct vec3_fp_vec_unit {
    fp_vec_unit_t x;
    fp_vec_unit_t y;
    fp_vec_unit_t z;
};
struct vec3_fp_color_channel {
    fp_color_channel_t r;
    fp_color_channel_t g;
    fp_color_channel_t b;
};
struct vec2_fp_vec_unit {
    fp_vec_unit_t u;
    fp_vec_unit_t v;
};


// -----------------------------------------------------------
// 通用数学辅助函数 (模板化以支持不同的定点数类型)
// -----------------------------------------------------------
template<typename T>
T clamp(T val, T min_val, T max_val) {
    if (val < min_val) return min_val;
    if (val > max_val) return max_val;
    return val;
}

template<typename T>
T saturate(T val) {
    return clamp(val, (T)0.0, (T)1.0);
}

// 添加 clamp_int 函数
template<typename T>
T clamp_int(T val, T min_val, T max_val) {
    if (val < min_val) return min_val;
    if (val > max_val) return max_val;
    return val;
}

// -----------------------------------------------------------
// 2. 基础算术运算集合 (模板函数，支持不同的 ap_fixed 类型)
// -----------------------------------------------------------
template<int W, int I, ap_q_mode Q, ap_o_mode O, int N>
ap_fixed<W, I, Q, O, N> add(ap_fixed<W, I, Q, O, N> a, ap_fixed<W, I, Q, O, N> b) { 
    return a + b; 
}
template<int W, int I, ap_q_mode Q, ap_o_mode O, int N>
ap_fixed<W, I, Q, O, N> sub(ap_fixed<W, I, Q, O, N> a, ap_fixed<W, I, Q, O, N> b) { 
    return a - b; 
}
// Modified mul function to handle bit width expansion for intermediate products
template<int W, int I, ap_q_mode Q, ap_o_mode O, int N>
ap_fixed<W, I, Q, O, N> mul(ap_fixed<W, I, Q, O, N> a, ap_fixed<W, I, Q, O, N> b) { 
    typedef ap_fixed<2*W, 2*I, Q, O, N> intermediate_prod_t;
    return (ap_fixed<W, I, Q, O, N>)((intermediate_prod_t)a * (intermediate_prod_t)b); 
}

template<int W, int I, ap_q_mode Q, ap_o_mode O, int N>
ap_fixed<W, I, Q, O, N> div(ap_fixed<W, I, Q, O, N> a, ap_fixed<W, I, Q, O, N> b) {
    ap_fixed<W, I, Q, O, N> zero_val = (ap_fixed<W, I, Q, O, N>)0;
    
    // 使用 LSB 的倍数作为阈值，更 HLS 友好
    ap_fixed<W, I, Q, O, N> lsb_val;
    lsb_val = 0;             // 首先清零所有位
    lsb_val.range(0, 0) = 1; // 设置最低位为 1
    ap_fixed<W, I, Q, O, N> threshold = lsb_val * (ap_fixed<W,I,Q,O,N>)4; // 例如 4 * LSB

    if (b == zero_val || (b > -threshold && b < threshold)) {
        return zero_val;
    }
    return a / b;
}



// --- LUT 定义 (放置在命名空间内，可在整个 fpal 库中访问) ---
constexpr int INV_SQRT_LUT_N_BITS = 7; // Number of bits used for LUT index
constexpr int INV_SQRT_LUT_SIZE = 1 << INV_SQRT_LUT_N_BITS;

// --- INV_SQRT_COORD_LUT (ap_fixed<24,14>) ---
// Generated for mantissa M in [0.5, 1.0), F=10, LUT_N_BITS=7
const ap_uint<24> INV_SQRT_COORD_LUT[] = {
    0x0005A5, 0x0005A0, 0x00059A, 0x000595, 0x00058F, 0x00058A, 0x000585, 0x000580,
    0x00057A, 0x000575, 0x000570, 0x00056B, 0x000566, 0x000561, 0x00055D, 0x000558,
    0x000553, 0x00054E, 0x00054A, 0x000545, 0x000540, 0x00053C, 0x000538, 0x000533,
    0x00052F, 0x00052A, 0x000526, 0x000522, 0x00051E, 0x00051A, 0x000515, 0x000511,
    0x00050D, 0x000509, 0x000505, 0x000501, 0x0004FD, 0x0004FA, 0x0004F6, 0x0004F2,
    0x0004EE, 0x0004EA, 0x0004E7, 0x0004E3, 0x0004DF, 0x0004DC, 0x0004D8, 0x0004D5,
    0x0004D1, 0x0004CE, 0x0004CA, 0x0004C7, 0x0004C3, 0x0004C0, 0x0004BD, 0x0004B9,
    0x0004B6, 0x0004B3, 0x0004B0, 0x0004AD, 0x0004A9, 0x0004A6, 0x0004A3, 0x0004A0,
    0x00049D, 0x00049A, 0x000497, 0x000494, 0x000491, 0x00048E, 0x00048B, 0x000488,
    0x000485, 0x000482, 0x00047F, 0x00047D, 0x00047A, 0x000477, 0x000474, 0x000471,
    0x00046F, 0x00046C, 0x000469, 0x000467, 0x000464, 0x000461, 0x00045F, 0x00045C,
    0x00045A, 0x000457, 0x000454, 0x000452, 0x00044F, 0x00044D, 0x00044A, 0x000448,
    0x000445, 0x000443, 0x000441, 0x00043E, 0x00043C, 0x00043A, 0x000437, 0x000435,
    0x000433, 0x000430, 0x00042E, 0x00042C, 0x000429, 0x000427, 0x000425, 0x000423,
    0x000420, 0x00041E, 0x00041C, 0x00041A, 0x000418, 0x000416, 0x000414, 0x000411,
    0x00040F, 0x00040D, 0x00040B, 0x000409, 0x000407, 0x000405, 0x000403, 0x000401,
};

// --- INV_SQRT_VEC_UNIT_LUT (ap_fixed<16,5>) ---
// Generated for mantissa M in [0.5, 1.0), F=11, LUT_N_BITS=7
const ap_uint<16> INV_SQRT_VEC_UNIT_LUT[] = {
    0x0B4B, 0x0B3F, 0x0B34, 0x0B2A, 0x0B1F, 0x0B14, 0x0B09, 0x0AFF,
    0x0AF5, 0x0AEA, 0x0AE0, 0x0AD6, 0x0ACC, 0x0AC3, 0x0AB9, 0x0AAF,
    0x0AA6, 0x0A9D, 0x0A93, 0x0A8A, 0x0A81, 0x0A78, 0x0A6F, 0x0A66,
    0x0A5D, 0x0A55, 0x0A4C, 0x0A44, 0x0A3B, 0x0A33, 0x0A2B, 0x0A23,
    0x0A1A, 0x0A12, 0x0A0B, 0x0A03, 0x09FB, 0x09F3, 0x09EB, 0x09E4,
    0x09DC, 0x09D5, 0x09CE, 0x09C6, 0x09BF, 0x09B8, 0x09B1, 0x09A9,
    0x09A2, 0x099C, 0x0995, 0x098E, 0x0987, 0x0980, 0x097A, 0x0973,
    0x096C, 0x0966, 0x095F, 0x0959, 0x0953, 0x094C, 0x0946, 0x0940,
    0x093A, 0x0934, 0x092E, 0x0928, 0x0922, 0x091C, 0x0916, 0x0910,
    0x090A, 0x0904, 0x08FF, 0x08F9, 0x08F3, 0x08EE, 0x08E8, 0x08E3,
    0x08DD, 0x08D8, 0x08D3, 0x08CD, 0x08C8, 0x08C3, 0x08BD, 0x08B8,
    0x08B3, 0x08AE, 0x08A9, 0x08A4, 0x089F, 0x089A, 0x0895, 0x0890,
    0x088B, 0x0886, 0x0881, 0x087C, 0x0878, 0x0873, 0x086E, 0x086A,
    0x0865, 0x0860, 0x085C, 0x0857, 0x0853, 0x084E, 0x084A, 0x0845,
    0x0841, 0x083D, 0x0838, 0x0834, 0x0830, 0x082B, 0x0827, 0x0823,
    0x081F, 0x081B, 0x0816, 0x0812, 0x080E, 0x080A, 0x0806, 0x0802,
};


// 辅助模板结构体，用于从 ap_fixed 类型中提取整数位宽 I
template<typename T>
struct get_i_width;

template<int W, int I, ap_q_mode Q, ap_o_mode O, int N>
struct get_i_width<ap_fixed<W, I, Q, O, N>> {
    static constexpr int value = I;
};

// 辅助模板结构体，用于从 ap_fixed 类型中提取小数位宽 F
template<typename T>
struct get_f_width;

template<int W, int I, ap_q_mode Q, ap_o_mode O, int N>
struct get_f_width<ap_fixed<W, I, Q, O, N>> {
    static constexpr int value = W - I;
};

// 牛顿-拉夫逊迭代法计算倒数平方根 (1/sqrt(val))
template<typename T, int NUM_ITERATIONS>
T custom_inv_sqrt_iterative(T val) {
    #pragma HLS INLINE off // 保持此函数不内联，以便HLS创建共享硬件
    #pragma HLS PIPELINE II=1 // 保持流水线化

    T zero_val_orig_type = (T)0;

    // Determine bitwidths of the original type
    constexpr int F_val = get_f_width<T>::value;
    constexpr int I_val = get_i_width<T>::value;
    constexpr int W_val = T::width;

    // --- 定义内部迭代类型 (internal_iter_t) 的位宽 ---
    // 增加小数位宽 F_ITER，以保持中间乘法精度。
    constexpr int F_ITER_INCREMENT = 2;
    constexpr int F_ITER = F_val + F_ITER_INCREMENT;

    constexpr int Max_I_For_InvSqrt_MinVal_Coord = 6;
    constexpr int Max_I_For_InvSqrt_MinVal_VecUnit = 6;

    constexpr int Base_I_For_Iter = std::is_same<T, fp_coord_t>::value ? 
                                       std::max(I_val, Max_I_For_InvSqrt_MinVal_Coord) : 
                                       std::max(I_val, Max_I_For_InvSqrt_MinVal_VecUnit);

    constexpr int I_ITER_OPT = Base_I_For_Iter + 2;

    constexpr int W_ITER = I_ITER_OPT + F_ITER; // Total width

    typedef ap_fixed<W_ITER, I_ITER_OPT, AP_RND_CONV, AP_SAT> internal_iter_t; 

    internal_iter_t val_iter = (internal_iter_t)val;
    internal_iter_t zero_val_iter = (internal_iter_t)0;

    // 1. 处理非正数输入
    if (val_iter <= zero_val_iter) { // Check using the wider type
        return zero_val_orig_type; // Return zero of the original type
    }

    // 2. 获取 ap_fixed 的原始无符号位模式
    ap_uint<W_val> val_bits = (ap_uint<W_val>)val.range(W_val - 1, 0); 

    if (val_bits == 0) { // Still possible if original val was positive but extremely small (less than LSB)
        return zero_val_orig_type;
    }

    // 3. 归一化输入值 (mantissa) 并计算指数 (exponent)
    int msb_pos = W_val - 1 - val_bits.countLeadingZeros(); 
    
    // F_val - 1 表示小数部分最高位的位置 (如 0.5 对应 2^-1)
    // 归一化后，msb_pos 应该指向 F_val - 1，所以 exponent 是 msb_pos 和 (F_val - 1) 的差
    int exponent = msb_pos - (F_val - 1); 

    ap_uint<W_val> normalized_mantissa_u = val_bits;
    if (exponent > 0) { 
        normalized_mantissa_u >>= exponent;
    } else if (exponent < 0) { 
        normalized_mantissa_u <<= -exponent;
    }

    // 4. LUT 索引
    // normalized_mantissa_u 代表 M * (2^F_val)，其中 M 在 [0.5, 1.0) 之间。
    // 在固定点中，0.5 对应整数值 1 << (F_val - 1)。
    ap_uint<W_val> M_minus_half_scaled = normalized_mantissa_u - (1 << (F_val - 1));
    // 将 [0, 0.5) * 2^F_val 的范围映射到 [0, INV_SQRT_LUT_SIZE)
    // 比例因子为 (INV_SQRT_LUT_SIZE * 2) / (2^F_val)
    ap_uint<INV_SQRT_LUT_N_BITS> lut_idx = (ap_uint<INV_SQRT_LUT_N_BITS>)((M_minus_half_scaled * (INV_SQRT_LUT_SIZE * 2)) >> F_val);

    internal_iter_t x0_iter; // Declare x0_iter here

    // 5. 从 LUT 获取初始猜测 x0
    // LUT 存储的是 1/sqrt(M) 的位模式。需要先加载到原始类型，再转换为迭代类型。
    if constexpr (std::is_same<T, fp_coord_t>::value) {
        fp_coord_t temp_x0;
        temp_x0.range(W_val - 1, 0) = INV_SQRT_COORD_LUT[lut_idx]; 
        x0_iter = (internal_iter_t)temp_x0;
    } else if constexpr (std::is_same<T, fp_vec_unit_t>::value) {
        fp_vec_unit_t temp_x0;
        temp_x0.range(W_val - 1, 0) = INV_SQRT_VEC_UNIT_LUT[lut_idx];
        x0_iter = (internal_iter_t)temp_x0;
    } else {
        // Fallback for unsupported types (very rough estimate)
        if (val_iter >= (internal_iter_t)1.0) x0_iter = (internal_iter_t)0.8;
        else x0_iter = (internal_iter_t)2.0;
    }

    // 6. 调整 x0 的指数
    // 1/sqrt(val) = 1/sqrt(M * 2^exponent) = (1/sqrt(M)) * 2^(-exponent/2)
    // 所以结果需要右移 exponent/2 位
    int result_shift = exponent / 2; 

    if (result_shift > 0) {
        x0_iter >>= result_shift;
    } else if (result_shift < 0) {
        x0_iter <<= -result_shift; 
    }

    // 7. 牛顿-拉夫逊迭代
    internal_iter_t one_point_five = (internal_iter_t)1.5;
    internal_iter_t point_five = (internal_iter_t)0.5;

    for (int i = 0; i < NUM_ITERATIONS; ++i) {
        #pragma HLS LOOP_TRIPCOUNT min=NUM_ITERATIONS max=NUM_ITERATIONS
        
        // 注意：这里 ap_fixed 的乘法会自动进行位宽扩展，然后根据赋值目标类型进行截断。
        // 由于 val_mul_x0_sq 是 internal_iter_t，其小数位宽 F_ITER 已经比 F_val 大，
        // 故在两次乘法后，即使被截断回 internal_iter_t，仍能保持比原始 F_val 更高的精度。
        internal_iter_t val_mul_x0_sq = val_iter * x0_iter * x0_iter; 
        internal_iter_t term = one_point_five - (point_five * val_mul_x0_sq);
        x0_iter = x0_iter * term;
    }

    return (T)x0_iter; // Convert final result back to the original fixed-point type T
}


// 牛顿-拉夫逊迭代法计算平方根 (sqrt(val))
template<typename T, int NUM_ITERATIONS>
T custom_sqrt_iterative(T val) {
    T zero_val_orig_type = (T)0; // A zero of the original type
    if (val <= zero_val_orig_type) return zero_val_orig_type; // Use the original type zero for comparison
    
    T inv_s = custom_inv_sqrt_iterative<T, NUM_ITERATIONS>(val);
    T res = mul(val, inv_s); // Use fpal::mul() for consistency
    
    return res;
}


// 非模板函数声明 (供外部使用，指定类型和迭代次数)
fp_coord_t sqrt(fp_coord_t val);
fp_vec_unit_t sqrt(fp_vec_unit_t val);
fp_coord_t inv_sqrt(fp_coord_t val);
fp_vec_unit_t inv_sqrt(fp_vec_unit_t val);


// 四舍五入到最近整数的辅助函数 (用于颜色转换)
template<typename T>
fp_color_channel_t round_to_nearest_uint8(T val) {
    // 这里的逻辑实现的是 "round half up"（四舍五入，0.5 向上取整）
    T rounded_val_for_conversion = val + (T)0.5; 
    return (fp_color_channel_t)rounded_val_for_conversion.to_uint(); 
}


// -----------------------------------------------------------
// 3. 线性代数运算集合 - 前置声明
// -----------------------------------------------------------

// 向量加法
vec3_fp_coord vec3_add(vec3_fp_coord a, vec3_fp_coord b);
vec3_fp_vec_unit vec3_add(vec3_fp_vec_unit a, vec3_fp_vec_unit b);

// 向量减法
vec3_fp_coord vec3_sub(vec3_fp_coord a, vec3_fp_coord b);
vec3_fp_vec_unit vec3_sub(vec3_fp_vec_unit a, vec3_fp_vec_unit b);

// 点积
fp_coord_t vec3_dot(vec3_fp_coord a, vec3_fp_coord b);
fp_vec_unit_t vec3_dot(vec3_fp_vec_unit a, vec3_fp_vec_unit b);

// 叉积
vec3_fp_coord vec3_cross(vec3_fp_coord a, vec3_fp_coord b);
vec3_fp_vec_unit vec3_cross(vec3_fp_vec_unit a, vec3_fp_vec_unit b);

// 向量长度平方
fp_coord_t vec3_length_sq(vec3_fp_coord v);
fp_vec_unit_t vec3_length_sq(vec3_fp_vec_unit v);
fp_coord_t vec3_length_sq(fp_coord_t x, fp_coord_t y, fp_coord_t z);
fp_vec_unit_t vec3_length_sq(fp_vec_unit_t x, fp_vec_unit_t y, fp_vec_unit_t z);


// 向量长度
fp_coord_t vec3_length(vec3_fp_coord v);
fp_coord_t vec3_length(fp_coord_t x, fp_coord_t y, fp_coord_t z);
fp_vec_unit_t vec3_length(vec3_fp_vec_unit v);

// 向量归一化
vec3_fp_coord vec3_normalize(vec3_fp_coord v);
vec3_fp_vec_unit vec3_normalize(vec3_fp_vec_unit v);
vec3_fp_coord vec3_normalize(fp_coord_t x, fp_coord_t y, fp_coord_t z);
vec3_fp_vec_unit vec3_normalize(fp_vec_unit_t x, fp_vec_unit_t y, fp_vec_unit_t z);


// 矩阵-向量乘法
vec3_fp_coord mat4_vec4_mul(const fp_coord_t mat[4][4], vec3_fp_coord vec_xyz, fp_coord_t w_in);
vec3_fp_vec_unit mat3_vec3_mul(const fp_vec_unit_t mat[3][3], vec3_fp_vec_unit vec);

// -----------------------------------------------------------
// 4. 定点数类型转换函数 - 前置声明
// -----------------------------------------------------------
fp_vec_unit_t conv_coord_to_vec_unit(fp_coord_t val);
fp_coord_t conv_vec_unit_to_coord(fp_vec_unit_t val);

fp_color_channel_t conv_fp_to_color_channel(fp_vec_unit_t val);
fp_color_channel_t conv_fp_to_color_channel(fp_coord_t val);

} // namespace fpal

// -----------------------------------------------------------
// 5. HLS 顶层函数声明 (全局命名空间，用于HLS工具识别)
// -----------------------------------------------------------
void fpal_wrapper(
    fpal::fp_coord_t in_coord_x,
    fpal::fp_coord_t in_coord_y,
    fpal::fp_coord_t in_coord_z,
    fpal::fp_vec_unit_t in_vec_u,
    fpal::fp_vec_unit_t in_vec_v,
    fpal::fp_color_channel_t &out_color_r,
    fpal::fp_coord_t &out_coord_len_sq,
    fpal::fp_vec_unit_t &out_vec_norm_x
);


#endif // FPAL_HPP

